package net.cyue.ort.llm;

import ai.onnxruntime.*;
import net.cyue.ort.ModelSessionManager;
import net.cyue.ort.llm.cache.CacheManager;
import net.cyue.ort.llm.cache.OnnxCacheManager;
import net.cyue.ort.llm.config.ModelConfig;
import net.cyue.ort.llm.data.ModelInput;
import net.cyue.ort.llm.data.ModelOutput;

import java.util.*;

/**
 * ONNX模型适配器实现
 * 负责ONNX模型的推理操作
 */
public class OnnxModelAdapter implements ModelAdapter {
    private static final int BATCH_SIZE = 1;
    
    private final ModelSessionManager sessionManager;
    private final OrtEnvironment env;
    private final OrtSession session;
    private final ModelConfig config;
    private final CacheManager cacheManager;
    private final boolean needAttentionMask;
    private final boolean needPositionID;
    
    public OnnxModelAdapter(
        ModelSessionManager sessionManager,
        ModelConfig config,
        CacheManager cacheManager
    ) {
        this.sessionManager = sessionManager;
        this.env = sessionManager.getEnvironment();
        this.session = sessionManager.getSession();
        this.config = config;
        this.cacheManager = cacheManager;
        
        // 验证模型
        Set<String> inputNames = session.getInputNames();
        this.needAttentionMask = inputNames.contains("attention_mask");
        this.needPositionID = inputNames.contains("position_ids");
        
        if (!inputNames.contains("input_ids")) {
            throw new ModelException("Model missing required input nodes: input_ids");
        }
        if (session.getOutputNames().isEmpty()) {
            throw new ModelException("Model has no output nodes");
        }
    }
    
    @Override
    public ModelOutput runInference(ModelInput input) throws ModelException {
        try {
            List<Long> inputIds = input.getInputIds();
            if (inputIds == null || inputIds.isEmpty()) {
                throw new ModelException("Input IDs cannot be empty");
            }
            
            boolean isFirstRun = input.isFirstRun();
            int seqLen = isFirstRun ? inputIds.size() : 1;
            // 直接使用ModelInput中的pastSeqLen，而不是从缓存中推断
            int pastSeqLen = input.getPastSeqLen();
            
            // 如果是第一次运行，确保pastSeqLen为0
            if (isFirstRun) {
                pastSeqLen = 0;
            }
            
            Map<String, OnnxTensor> inputs = prepareInputs(
                inputIds,
                seqLen,
                pastSeqLen,
                input.getPastKeyValueCache(),
                isFirstRun
            );
            
            try {
                OrtSession.Result result = session.run(inputs);
                return parseResult(result);
            } finally {
                cleanupInputs(inputs);
            }
        } catch (OrtException e) {
            throw new ModelException("Inference failed", e);
        }
    }
    
    @SuppressWarnings("unchecked")
    private Map<String, OnnxTensor> prepareInputs(
        List<Long> inputIDList,
        int seqLen,
        int pastSeqLen,
        Map<String, Object> cacheMap,
        boolean isFirstRun
    ) {
        // 与 LLM.java 保持一致：如果是第一次运行，重新设置 seqLen 和 pastSeqLen
        if (isFirstRun) {
            seqLen = inputIDList.size();
            pastSeqLen = 0;
        }
        
        int totalSeqLen = config.isUseCache() ? seqLen + pastSeqLen : seqLen;
        
        // 创建数组
        long[][] inputIDArray = new long[BATCH_SIZE][seqLen];
        long[][] attentionMaskArray = new long[BATCH_SIZE][totalSeqLen + 1];
        long[][] positionIDArray = needPositionID ? new long[BATCH_SIZE][seqLen] : new long[0][0];
        
        for (int i = 0; i < BATCH_SIZE; i++) {
            // 填充 input ids
            for (int j = 0; j < seqLen; j++) {
                inputIDArray[i][j] = inputIDList.get(inputIDList.size() - seqLen + j);
                if (needPositionID && positionIDArray.length > 0) {
                    positionIDArray[i][j] = config.isUseCache() ? pastSeqLen + j : j;
                }
            }
            // 填充 attention_mask
            if (needAttentionMask) {
                Arrays.fill(attentionMaskArray[i], 1);
            }
        }
        
        Map<String, OnnxTensor> inputs = new HashMap<>();
        try {
            inputs.put("input_ids", OnnxTensor.createTensor(env, inputIDArray));
            if (needAttentionMask) {
                inputs.put("attention_mask", OnnxTensor.createTensor(env, attentionMaskArray));
            }
            if (needPositionID && positionIDArray.length > 0) {
                inputs.put("position_ids", OnnxTensor.createTensor(env, positionIDArray));
            }
            
            // 处理 past_key_values
            Map<String, Object> finalCache = null;
            if (config.isUseCache()) {
                @SuppressWarnings("unchecked")
                Map<String, OnnxTensor> onnxCache = cacheMap != null ? 
                    (Map<String, OnnxTensor>)(Map<?, ?>) cacheMap : null;
                
                if (onnxCache == null || onnxCache.isEmpty()) {
                    finalCache = cacheManager.createEmptyCache(
                        BATCH_SIZE,
                        pastSeqLen + 1,
                        config
                    );
                } else {
                    OnnxTensor keyTensor = onnxCache.get("past_key_values.0.key");
                    if (keyTensor != null && keyTensor.isClosed()) {
                        throw new ModelException("past_key_values cache is closed");
                    }
                    if (keyTensor != null) {
                        int cacheSeqLen = (int) keyTensor.getInfo().getShape()[2];
                        int expandLength = (pastSeqLen + 1) - cacheSeqLen;
                        if (expandLength > 0) {
                            finalCache = cacheManager.expandCache(cacheMap, expandLength, config);
                        } else {
                            finalCache = cacheMap;
                        }
                    }
                }
                
                if (finalCache != null) {
                    @SuppressWarnings("unchecked")
                    Map<String, OnnxTensor> tensorCache = (Map<String, OnnxTensor>)(Map<?, ?>) finalCache;
                    inputs.putAll(tensorCache);
                }
            } else if (session.getInputNames().contains("past_key_values.0.key")) {
                finalCache = cacheManager.createEmptyCache(BATCH_SIZE, 1, config);
                @SuppressWarnings("unchecked")
                Map<String, OnnxTensor> tensorCache = (Map<String, OnnxTensor>)(Map<?, ?>) finalCache;
                inputs.putAll(tensorCache);
            }
            
            return inputs;
        } catch (OrtException e) {
            cleanupInputs(inputs);
            throw new ModelException("Error creating input tensors", e);
        }
    }
    
    @SuppressWarnings("unchecked")
    private ModelOutput parseResult(OrtSession.Result result) throws OrtException {
        if (result == null) {
            throw new ModelException("Model output is null");
        }
        
        OnnxValue firstOutput = result.get(0);
        if (firstOutput == null) {
            throw new ModelException("Model output is empty");
        }
        
        try {
            float[][][] logits = (float[][][]) firstOutput.getValue();
            
            if (logits == null || logits.length == 0 || 
                logits[0] == null || logits[0].length == 0) {
                throw new ModelException("Logits array dimension error");
            }
            
            // 提取缓存
            Map<String, Object> cache = null;
            if (cacheManager instanceof OnnxCacheManager) {
                cache = ((OnnxCacheManager) cacheManager)
                    .extractFromOrtResult(result, config);
            }
            
            return new ModelOutput(logits, cache);
        } catch (ClassCastException e) {
            throw new ModelException("Model output type error: " + e.getMessage(), e);
        }
    }
    
    @SuppressWarnings("unchecked")
    private void cleanupInputs(Map<String, OnnxTensor> inputs) {
        if (inputs == null) return;
        inputs.values().stream()
            .filter(tensor -> tensor != null && !tensor.isClosed())
            .forEach(OnnxValue::close);
        inputs.clear();
    }
    
    @Override
    public ModelConfig getConfig() {
        return config;
    }
    
    @Override
    public void close() {
        sessionManager.close();
    }
}

